from openai import OpenAI
import configparser
from fastapi.responses import StreamingResponse
import asyncio

# 将 call_model 函数改为异步函数，并使其能逐步返回数据
async def call_model(message, mode):
    # 读取 config/model_config.ini 模型配置
    model_config = configparser.ConfigParser()
    model_config.read("config/model_config.ini")

    # 根据传入的参数，读取模型配置
    config = model_config[mode]

    # 初始化OpenAI客户端
    client = OpenAI(
        api_key=config["api_key"],
        base_url=config["host"]
    )

    # 创建聊天完成请求
    completion = client.chat.completions.create(
        model=config["mode"],
        messages=[
            {"role": "user", "content": message}
        ],
        stream=True,
        # 解除以下注释会在最后一个chunk返回Token使用量
        # stream_options={
        #     "include_usage": True
        # }
    )

    reasoning_content = ""  # 定义完整思考过程
    answer_content = ""     # 定义完整回复
    is_answering = False   # 判断是否结束思考过程并开始回复

    print("\n" + "=" * 20 + "思考过程" + "=" * 20 + "\n")

    async def generate():
        nonlocal reasoning_content, answer_content, is_answering  # 声明使用外部作用域的变量
        for chunk in completion:
            # 如果chunk.choices为空，则打印usage
            if not chunk.choices:
                print("\nUsage:")
                print(chunk.usage)
            else:
                delta = chunk.choices[0].delta
                # 打印思考过程
                if hasattr(delta, 'reasoning_content') and delta.reasoning_content is not None:
                    print(delta.reasoning_content, end='', flush=True)
                    reasoning_content += delta.reasoning_content
                else:
                    # 开始回复
                    if delta.content != "" and not is_answering:
                        print("\n" + "=" * 20 + "完整回复" + "=" * 20 + "\n")
                        is_answering = True
                    # 打印回复过程
                    print(delta.content, end='', flush=True)
                    answer_content += delta.content
                # 每次处理完一个chunk，就将当前的结果返回一部分
                partial_result = {
                    "message": reasoning_content + answer_content,
                    "products": []
                }
                json_str = str(partial_result).replace("'", "\"")
                yield f"data: {json_str}\n\n"
        # 处理完所有chunk后，返回最终结果
        final_result = {
            "message": reasoning_content + answer_content,
            "products": []
        }
        json_str = str(final_result).replace("'", "\"")
        yield f"data: {json_str}\n\n"

    return generate()